from airflow import DAG

from datetime import datetime, timedelta

from airflow.contrib.operators.kubernetes_pod_operator import KubernetesPodOperator
from airflow.kubernetes.secret import Secret
from airflow.operators.dummy_operator import DummyOperator
from kubernetes.client import models as k8s

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2024, 7, 31),
    'start_date': datetime.utcnow(),
    'email': ['airflow@example.com'],
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5)
}


volume = k8s.V1Volume(
    name="cache-volume",
    host_path={
      "path": "/data/v8/airflow_temp", # directory location on host
      "type": "DirectoryOrCreate"} # this field is optional
)

volume_mounts = [
    k8s.V1VolumeMount(
        mount_path="/cache", name="cache-volume"
    )
]


container_resources_cuda_gpu = {
    "request_memory": "55G",
    "request_cpu": "15000m",
    "limit_gpu": "1",
}

#   requests:
#     cpu: 250m
#     memory: 512Mi
#     nvidia.com/gpu: 1 # 申请GPU的数量
#   limits:
#     cpu: 250m
#     memory: 512Mi
#     nvidia.com/gpu: 1 # GPU数量的使用上限


container_resources_cuda_gpu=k8s.V1ResourceRequirements(
        requests={"cpu": "10000m", "memory": "6400Mi", "nvidia.com/gpu": "1"},
        limits={"cpu": "20000m", "memory": "6400Mi", "nvidia.com/gpu": "1"},
    )


dag = DAG('kubernetes_sample_gpu', default_args=default_args, schedule_interval=timedelta(days=1))



passing = KubernetesPodOperator(namespace='airflow',
                          image="k8s-pod-test-gpu:2.0",
                          labels={"foo": "bar"},
                          name="passing-test",
                          task_id="passing-task",
                          get_logs=True,
                          is_delete_operator_pod=True,
                          in_cluster=True,
                          env_vars= {"test_secret": "value123"},
                          volume_mounts=volume_mounts,
                          volumes=[volume],
                          #resources=container_resources_cuda_gpu,
                          dag=dag
                          )



